from bs4 import BeautifulSoup
import re
from collections import Counter

# 1. 使用Beautiful Soup构建TEI文件的树结构
with open('oth.xml', 'r', encoding='utf-8') as file:
    content = file.read()
soup = BeautifulSoup(content, 'xml')

# 2. 从TEI头部提取标题
title = soup.find('title').text
print(f"Extracted Title: {title}")

# 3. 提取body中的所有文本
body_text = []
for w_tag in soup.find_all('w'):
    body_text.append(w_tag.get_text())
body_text_str = ' '.join(body_text)
print(f"Extracted Body Text: {body_text_str[:200]}...")  # 打印前200个字符

# 4. 创建一个新的纯文本文件
title_cleaned = re.sub(r'\W+', '_', title.lower()).strip('_')  # 清理标题以用于文件名
output_file_path = f'{title_cleaned}.txt'
# output_filename = 'othello_the_moor_of_venice.txt'
with open(output_file_path, 'w', encoding='utf-8') as output_file:
    output_file.write(body_text_str)
print(f"Text saved to {output_file_path}")

# 5. 提取Iago的所有发言，并统计每个单词的出现次数
def extract_speech(speaker_name):
    speeches = soup.find_all('sp', {'who': speaker_name})
    speech_text = ' '.join([speech.get_text(separator=' ', strip=True) for speech in speeches])
    words = re.findall(r'\b\w+\b', speech_text.lower())
    word_count = Counter(words)
    return word_count

iago_word_count = extract_speech('Iago')

# 创建一个新文件，用于存储Iago的单词计数
iago_output_filename = 'iago_word_count.txt'
with open(iago_output_filename, 'w', encoding='utf-8') as iago_output_file:
    for word, count in iago_word_count.items():
        iago_output_file.write(f"{word}:{count}\n")
print(f"Iago's word count saved to {iago_output_filename}")

# 6. 创建一个函数，接受一个字符串作为参数，并返回一个单词列表
# def extract_words(text):
#     words = re.findall(r'\b\w+\b', text.lower())  # todo 用replace
#     return words

def extract_words(text):
    # 使用正则表达式替换标点符号为空格
    text = re.sub(r'[.,!?;:"\'-()[]{}<>]', ' ', text)
    # 将文本转换为小写并分割成单词列表
    words = text.lower().split()
    return words

# 7. 统计每幕的总单词数，并找出哪一幕的单词数最多 todo 还需要确认如何提取每一幕的结构化
# 提取每幕的文本
act_pattern = r'<div xml:id="([^"]+)">\s*<head>Act \d+</head>.*?</div>'
acts = re.findall(act_pattern, content, re.DOTALL)
act_word_counts = {}

for act_text in acts:
    words = re.findall(r'\b\w+\b', act_text.lower())
    act_word_counts[act_text] = len(words)

# 程序化地确定哪一幕的单词数最多
max_act = max(act_word_counts, key=act_word_counts.get)
max_word_count = act_word_counts[max_act]

print(f"Act with the most words: {max_act} ({max_word_count} words)")

# 将每幕的单词数存储到文件中
act_word_count_filename = 'act_word_counts.txt'
with open(act_word_count_filename, 'w', encoding='utf-8') as act_word_count_file:
    for act_id, word_count in act_word_counts.items():
        act_word_count_file.write(f"{act_id}: {word_count} words\n")
print(f"Act word counts saved to {act_word_count_filename}")